Cache pre-existing Zarr arrays in Zarr backend #9861

d-v-b · 2024-12-06T13:43:25Z

ZarrStore calls array_keys on the zarr.Group it wraps multiple times. Each of these array_keys calls performs IO, but if we model the names of the arrays as static, then they can be cached, then we can do less IO.

Accordingly this PR adds a new method to ZarrStore (get_array_keys), a new keyword argument for ZarrStore.__init__ (cache_array_keys), and a new attribute to ZarrStore (_cache). get_array_keys gets the array keys from the .zarr_group instance contained in the ZarrStore, either by reading the keys from the cache (if caching is enabled and the cache is populated), or by directly invoking ZarrStore.zarr_group.array_keys() (if caching is disabled, or the cache is empty).

Caching is controlled by the cache_array_keys keyword argument in ZarrStore.__init__.

If ZarrStore.__init__ is called with with cache_array_keys=True, the _cache attribute of the ZarrStore is set to {"array_keys": None}. On the first invocation of get_array_keys, this empty-but-active cache is filled with a call to zarr_group.array_keys, and that value is returned. On subsequent invocations of get_array_keys, the value from the cache is retrieved and returned.
If ZarrStore.__init__ is called with cache_array_keys=False, then the _cache attribute never gets a array_keys key inserted, and this indicates to get_array_keys that the cache is not being used. Calls to get_array_keys will return the result of zarr_group.array_keys.

Let me know if the _cache thing is a bit baroque -- my initial effort added a _cache_array_keys attribute as well as the cache, but I figured the state of the cache itself can convey whether the cache is meant to be used, which removes the need for a "use the cache or not" attribute on ZarrStore.

And I would like some feedback on the test -- is it in the right place?

closes #9853

welcome · 2024-12-06T13:43:28Z

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

for more information, see https://pre-commit.ci

… into perf/cache-array-keys

…/cache-array-keys

dcherian · 2024-12-06T15:01:32Z

If you swtich the equality to != we should see a decrease in list_prefix:

xarray/xarray/tests/test_backends.py

Line 3325 in eac5105

assert summary[k] <= expected[k], (k, summary)

xarray/backends/zarr.py

d-v-b · 2024-12-06T17:23:29Z

@dcherian I have implemented caching both array keys and arrays, but for that reason this test is failing:

xarray/xarray/tests/test_backends.py

Lines 409 to 416 in eac5105

    
           def test_write_store(self) -> None: 
        
               expected = create_test_data() 
        
               with self.create_store() as store: 
        
                   expected.dump_to_store(store) 
        
                   # we need to cf decode the store because it has time and 
        
                   # non-dimension coordinates 
        
                   with xr.decode_cf(store) as actual: 
        
                       assert_allclose(expected, actual)

This test assumes that the ZarrStore will track changes to the underlying Zarr data which makes caching tricky. Does re-using the ZarrStore here reflect real-world usage, or is it just for testing convenience? and thus should I modify this test to create a new ZarrStore, or explicitly turn off the cache for this test?

dcherian · 2024-12-06T17:31:23Z

oh wow, I did not even know that there was a Dataset.dump_to_store method!

This code is from 10 years ago and I don't think anyone is using it like this. Today this test would be

with self.roundtrip(expected) as actual:
    assert_identical(actual, expected)

cc @shoyer

xarray/backends/zarr.py

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

…he-array-keys

d-v-b · 2024-12-10T09:17:40Z

oh wow, I did not even know that there was a Dataset.dump_to_store method!

This code is from 10 years ago and I don't think anyone is using it like this. Today this test would be
with self.roundtrip(expected) as actual:
    assert_identical(actual, expected)
cc @shoyer

I altered the test as you recommended, let me know if that change should be in a separate PR

… into perf/cache-array-keys

xarray/backends/zarr.py

xarray/tests/test_backends.py

xarray/backends/zarr.py

xarray/tests/test_backends.py

dcherian · 2024-12-12T03:36:43Z

xarray/tests/test_backends.py

            yield group

+    def test_write_store(self) -> None:


does this need to be over-ridden still if the default is to turn the cache off

I think so, because we need to explicitly ensure that the cache is off, otherwise tests will fail. If we rely on the default behavior (which is defined very far away from this test), and if that default value changes, then tests will fail here and it will not be clear why.

another way to remove this code would be to change the tests so that they no longer rely on the old ZarrStore behavior; maybe this could be a separate PR.

I'll push a commit. The create_store on this class seems pointless.

dcherian

Very nice! Clearly its pretty effective. I left some minor suggestions.

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

jhamman

Looks great @d-v-b. Left a comment on the default value.

Also, please add a what's new note in the docs.

jhamman · 2024-12-14T22:46:24Z

xarray/backends/zarr.py

@@ -633,6 +635,7 @@ def open_store(
        zarr_format=None,
        use_zarr_fill_value_as_mask=None,
        write_empty: bool | None = None,
+        cache_members: bool = False,


What do you think about changing the default to True? I'd like to know if we think there is a downside (risk, performance, etc.) to defaulting to this behavior?

There are some weird tests that use open_store and dump_to_store that can't handle that default.

For nearly all users, the effective default is the value set in open_zarr and to_zarr (which is True, or should be after that one suggestion is merge)

i'm up for it! the main risk is whether anyone is actively using ZarrStore instances explicitly, with the assumption that they dynamically track the state of the underlying zarr Group. if nobody is doing that, then there's no risk. ofc I have no way of know if this use case is common, but my impression was that ZarrStore is internal enough that users are not expected to be using it directly.

xarray/backends/zarr.py

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

…/cache-array-keys

d-v-b · 2024-12-17T12:55:57Z

CI is failing on mypy due to the following:

 Found 5 errors in 1 file (checked 170 source files)
xarray/tests/test_plot.py:98: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:105: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:113: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:125: error: Unused "type: ignore" comment  [unused-ignore]

I don't think that's from this PR?

d-v-b · 2024-12-17T14:15:59Z

Looks great @d-v-b. Left a comment on the default value.

Also, please add a what's new note in the docs.

"what's new" has been added and the default value is now True. CI is still red due to some mypy issues that I think are unrelated to my pr (see #9861 (comment))

welcome · 2024-12-18T16:19:43Z

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again!

d-v-b added 3 commits December 4, 2024 21:34

cache array keys on ZarrStore instances

9503c71

refactor get_array_keys

94b48ac

add test for get_array_keys

2dacb2a

pre-commit-ci bot and others added 5 commits December 6, 2024 13:45

[pre-commit.ci] auto fixes from pre-commit.com hooks

c7cfacf

for more information, see https://pre-commit.ci

lint

b3cf2e5

Merge branch 'perf/cache-array-keys' of https://github.com/d-v-b/xarray…

d9bd7fd

… into perf/cache-array-keys

Merge branch 'main' of https://www.github.com/pydata/xarray into perf…

8d47dd3

…/cache-array-keys

add type annotation for _cache

bec8b42

d-v-b added 3 commits December 6, 2024 16:02

make type hint accurate

61fe759

make type hint pass mypy

0952c93

make type hint pass mypy, correctly

96ddcf4

dcherian reviewed Dec 6, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian reviewed Dec 9, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

d-v-b and others added 3 commits December 9, 2024 21:19

Update xarray/backends/zarr.py

60761a8

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Merge branch 'main' of https://github.com/pydata/xarray into perf/cac…

f82659d

…he-array-keys

use roundtrip method instead of explicit alternative

107c4e7

d-v-b added 4 commits December 10, 2024 10:24

cache members instead of just array keys

2d3c13a

adjust test to members caching

be5e389

Merge branch 'perf/cache-array-keys' of https://github.com/d-v-b/xarray…

1eaf3ea

… into perf/cache-array-keys

refactor members cache

9d147b9

dcherian reviewed Dec 10, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian reviewed Dec 10, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian reviewed Dec 10, 2024

View reviewed changes

xarray/tests/test_backends.py Outdated Show resolved Hide resolved

d-v-b force-pushed the perf/cache-array-keys branch 2 times, most recently from a7602f6 to 9d147b9 Compare December 10, 2024 16:57

dcherian reviewed Dec 12, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian reviewed Dec 12, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian reviewed Dec 12, 2024

View reviewed changes

xarray/tests/test_backends.py Show resolved Hide resolved

dcherian reviewed Dec 12, 2024

View reviewed changes

xarray/tests/test_backends.py Show resolved Hide resolved

dcherian reviewed Dec 12, 2024

View reviewed changes

dcherian approved these changes Dec 12, 2024

View reviewed changes

dcherian requested a review from jhamman December 12, 2024 03:37

d-v-b and others added 4 commits December 14, 2024 12:26

Update xarray/backends/zarr.py

d987529

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Update xarray/backends/zarr.py

169b736

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Update xarray/backends/zarr.py

e20d11d

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Update xarray/backends/zarr.py

a1c25f2

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

jhamman reviewed Dec 14, 2024

View reviewed changes

dcherian reviewed Dec 16, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

d-v-b and others added 6 commits December 16, 2024 17:50

Update xarray/backends/zarr.py

f8d78cb

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

doc: add whats new content

650937e

fixup

fcf2d64

Merge branch 'main' of https://www.github.com/pydata/xarray into perf…

bf64fd2

…/cache-array-keys

update whats new

dcecf43

fixup

d0bb5a2

d-v-b added 2 commits December 17, 2024 14:05

remove vestigial cache_array_keys

65cadcb

make cache_members default to True

e8cd3c8

dcherian added 2 commits December 17, 2024 10:00

tweak

8d31f0e

Remove from public API

3d15d3d

dcherian added the plan to merge Final call for comments label Dec 17, 2024

Merge branch 'main' into perf/cache-array-keys

f1ef905

dcherian enabled auto-merge (squash) December 18, 2024 15:48

dcherian merged commit 60a4a49 into pydata:main Dec 18, 2024
29 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Cache pre-existing Zarr arrays in Zarr backend #9861

Cache pre-existing Zarr arrays in Zarr backend #9861

d-v-b commented Dec 6, 2024 •

edited

Loading

welcome bot commented Dec 6, 2024

dcherian commented Dec 6, 2024

d-v-b commented Dec 6, 2024

dcherian commented Dec 6, 2024

d-v-b commented Dec 10, 2024

dcherian Dec 12, 2024

d-v-b Dec 17, 2024

d-v-b Dec 17, 2024

dcherian Dec 17, 2024

dcherian left a comment

jhamman left a comment

jhamman Dec 14, 2024

dcherian Dec 16, 2024

d-v-b Dec 16, 2024

d-v-b commented Dec 17, 2024

d-v-b commented Dec 17, 2024

welcome bot commented Dec 18, 2024

Cache pre-existing Zarr arrays in Zarr backend #9861

Cache pre-existing Zarr arrays in Zarr backend #9861

Conversation

d-v-b commented Dec 6, 2024 • edited Loading

welcome bot commented Dec 6, 2024

dcherian commented Dec 6, 2024

d-v-b commented Dec 6, 2024

dcherian commented Dec 6, 2024

d-v-b commented Dec 10, 2024

dcherian Dec 12, 2024

Choose a reason for hiding this comment

d-v-b Dec 17, 2024

Choose a reason for hiding this comment

d-v-b Dec 17, 2024

Choose a reason for hiding this comment

dcherian Dec 17, 2024

Choose a reason for hiding this comment

dcherian left a comment

Choose a reason for hiding this comment

jhamman left a comment

Choose a reason for hiding this comment

jhamman Dec 14, 2024

Choose a reason for hiding this comment

dcherian Dec 16, 2024

Choose a reason for hiding this comment

d-v-b Dec 16, 2024

Choose a reason for hiding this comment

d-v-b commented Dec 17, 2024

d-v-b commented Dec 17, 2024

welcome bot commented Dec 18, 2024

d-v-b commented Dec 6, 2024 •

edited

Loading